"""
Classification demo for Fashion MNIST using the NeuralNetwork class.

To run this demo, use:
    python -m demos.fashion_mnist.demo-fashion-mnist
from the project root directory.
"""
import numpy as np
import gzip
import matplotlib.pyplot as plt
import os

# Import the neural network components from the new package
from neural_network import Layer
from neural_network import ReLUActivation
from neural_network import SigmoidActivation
from neural_network import NeuralNetwork
from neural_network import LossSoftmaxCCE

# Settings
hidden_sizes = [128, 128, 128]
# hidden_sizes = [1000]*4
# hidden_sizes = [4, 4]
activation_hidden = SigmoidActivation
epochs = 1000
learning_rate = 0.01
batch_size = 64

# Class labels
label_names = [
    "T-shirt/top", "Trouser", "Pullover", "Dress", "Coat",
    "Sandal", "Shirt", "Sneaker", "Bag", "Ankle boot"
]

print("Loading Fashion MNIST dataset from local CSV files...")
# Use package-relative paths to the data files
data_dir = os.path.dirname(__file__)
train_path = os.path.join(data_dir, 'fashion-mnist_train.csv.gz')
test_path = os.path.join(data_dir, 'fashion-mnist_test.csv.gz')
train_data = np.loadtxt(gzip.open(train_path, 'rt'), delimiter=',', skiprows=1)
test_data = np.loadtxt(gzip.open(test_path, 'rt'), delimiter=',', skiprows=1)

# Split features and labels
y_raw_train = train_data[:, 0].astype(int)
X_train = train_data[:, 1:].T / 255.0
y_raw_test = test_data[:, 0].astype(int)
X_test = test_data[:, 1:].T / 255.0

num_classes = len(label_names)
y_train = np.eye(num_classes)[y_raw_train].T
y_test = np.eye(num_classes)[y_raw_test].T

# Build the network
input_size = X_train.shape[0]
network = NeuralNetwork()
prev_size = input_size
for size in hidden_sizes:
    network.add_layer(Layer(prev_size, size))
    network.add_layer(activation_hidden())
    prev_size = size
network.add_layer(Layer(prev_size, num_classes))

# Loss function
loss_obj = LossSoftmaxCCE()

# Use the enhanced train method with accuracy metric
print("Training the neural network...")
network.train(
    X_train, y_train,
    loss_obj,
    epochs,
    learning_rate,
    batch_size,
    shuffle=True,
    verbose=True,
    X_val=X_test,
    y_val=y_test,
    metric="accuracy"
)

# Get final predictions for later use
y_pred_test = network.forward_pass(X_test)

print("Training complete. Showing test images one by one. Close the window to advance.")

# Display each test image with predicted logits and probabilities
for i in range(X_test.shape[1]):
    x = X_test[:, i].reshape(-1, 1)
    logits = network.forward_pass(x)
    exp_vals = np.exp(logits - np.max(logits, axis=0, keepdims=True))
    probs = exp_vals / np.sum(exp_vals, axis=0, keepdims=True)

    fig, axes = plt.subplots(1, 2, figsize=(8, 4))
    # Show image
    axes[0].imshow(x.reshape(28, 28), cmap='gray')
    axes[0].set_title(f"True: {label_names[y_raw_test[i]]}")
    axes[0].axis('off')

    # Bar plot of probabilities
    axes[1].bar(range(num_classes), probs.flatten(), color='skyblue')
    axes[1].set_xticks(range(num_classes))
    axes[1].set_xticklabels(label_names, rotation=45, ha='right')
    axes[1].set_ylim(0, 1)
    axes[1].set_title("Predicted probabilities")

    plt.tight_layout()
    plt.show()